package me.linx.demo

import at.favre.lib.crypto.HKDF
import org.junit.Test
import java.nio.ByteBuffer
import java.security.MessageDigest
import java.security.SecureRandom
import java.util.*
import javax.crypto.Cipher
import javax.crypto.Mac
import javax.crypto.SecretKey
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec

class EncryptTest {

    @Test
    fun encryptTest() {
        val util = EncryptionUtil()
        val result = util.encryptBasic(
            UUID.fromString("20e086a6-1961-4d53-ba61-695a66b0351f"),
            "A01AB".toByteArray()
        )
        println(result)
    }


}

class EncryptionUtil {

    companion object {
        const val KEY_BYTE_SIZE = 16
        const val ENC_KEY_INFO = "BRBTCS_STARK_ENC_KEY"
        const val AUTH_KEY_INFO = "BRBTCS_STARK_AUTH_KEY"
        const val FINAL_KEY_INFO = "BRBTCS_STARK_FINAL_KEY"
    }

    /**
     * Derive an encryption key using HKDF.
     *
     * @param source source material
     * @param info context information
     */
    fun deriveKey128(source: ByteArray, info: String): ByteArray {
        return HKDF.fromHmacSha256().expand(source, info.toByteArray(), 16)
    }

    /**
     * AES-128-CBC encryption with HMAC-SHA256 signature.
     *
     * @param key 128-bit AES key
     * @param plainText plain text
     */
    fun encryptAdvanced(key: String, plainText: ByteArray): ByteArray {
        val secureRandom = SecureRandom()
        val keyBytes = key.toByteArray()
        // Make sure key size is KEY_BYTE_SIZE
        if (keyBytes.size != KEY_BYTE_SIZE) {
            throw IllegalArgumentException(
                "Key length size illegal: " +
                        "${KEY_BYTE_SIZE * 8}-bit is expected but ${keyBytes.size}-bit key is detected"
            )
        }

        // Generate initial vector
        val iv = ByteArray(16)
        secureRandom.nextBytes(iv)

        // Generate encKey and authKey
        val encKey = HKDF.fromHmacSha256().expand(keyBytes, ENC_KEY_INFO.toByteArray(), 16)
        val authKey = HKDF.fromHmacSha256()
            .expand(keyBytes, AUTH_KEY_INFO.toByteArray(), 32) // //HMAC-SHA256 key is 32 byte

        // Compute AES encrypted cipherText
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding") //actually uses PKCS#7
        cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(encKey, "AES"), IvParameterSpec(iv))
        val cipherText = cipher.doFinal(plainText)

        // Compute HmacSHA256 Signature
        val macKey: SecretKey = SecretKeySpec(authKey, "HmacSHA256")
        val hmac = Mac.getInstance("HmacSHA256")
        hmac.init(macKey) // Initialize mac key
        hmac.update(iv)
        hmac.update(cipherText)

        val mac = hmac.doFinal()

        // Compute cipher message, combining iv, mac, cipherText to cipherMessage
        val byteBuffer = ByteBuffer.allocate(1 + iv.size + 1 + mac.size + cipherText.size)
        byteBuffer.put(iv.size.toByte())
        byteBuffer.put(iv)
        byteBuffer.put(mac.size.toByte())
        byteBuffer.put(mac)
        byteBuffer.put(cipherText)

        return byteBuffer.array()
    }

    fun encryptBasic(userId: UUID, plainText: ByteArray): String {
        val secureRandom = SecureRandom()
        // Generate initial vector
        val iv = byteArrayOf(
            -62, -12, 85, 80, -34, 58, -91, 31, 64, 101, -39, 65, -88, 53, -16, -88
        )
//        secureRandom.nextBytes(iv)
        println(iv.toList())
        println(userId.toString())
        val key = getMD5(userId)
        println("key ${key.toList()}")
        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding") //actually uses PKCS#7
        cipher.init(Cipher.ENCRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
        val cipherText = cipher.doFinal(plainText)
        println("cipherText ${cipherText.toList()}")
        // Compute HmacSHA256 Signature
        val macKey: SecretKey = SecretKeySpec(key, "HmacSHA256")
        val hmac = Mac.getInstance("HmacSHA256")
        hmac.init(macKey) // Initialize mac key
        hmac.update(iv)
        hmac.update(cipherText)

        val mac = hmac.doFinal()
        println("mac ${mac.toList()}")
        val byteBuffer = ByteBuffer.allocate(iv.size + mac.size + cipherText.size)
        byteBuffer.put(iv)
        byteBuffer.put(cipherText)
        byteBuffer.put(mac)
        println("final buffer ${byteBuffer.array().toList()}")
        return Base64.getEncoder().encodeToString(byteBuffer.array())
    }

    private fun getBytesFromUUID(id: UUID): ByteArray {
        println("${id.mostSignificantBits} ${id.leastSignificantBits}")
        val bb = ByteBuffer.wrap(ByteArray(16))
        bb.putLong(id.mostSignificantBits)
        bb.putLong(id.leastSignificantBits)
        return bb.array()
    }

    private fun getMD5(id: UUID): ByteArray {
        val md = MessageDigest.getInstance("MD5")
        val bytes = getBytesFromUUID(id)
        println("bytes ${bytes.toList()}")
        md.update(getBytesFromUUID(id))
        return md.digest()
    }

    fun decryptBasic(userId: UUID, encStr: String): String {
        val encText = Base64.getDecoder().decode(encStr)
        val payloadSize = encText.size - 32 - 16
        val encrypted = encText.copyOfRange(16, 16 + payloadSize)
        val iv = encText.copyOfRange(0, 16)

        val key = getMD5(userId)

        val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding") //actually uses PKCS#7
        cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
        val decrypted = cipher.doFinal(encrypted)
        return decrypted.toString(Charsets.UTF_8)
    }
}